/********************************************************************
	created:	30:9:2007   22:34
	author:		tiamo
	purpose:	helper
*********************************************************************/

#include "stdafx.h"

#ifdef ALLOC_PRAGMA
	#pragma alloc_text(PAGE,AppleKeyboardCallHidClass)
	#pragma alloc_text(PAGE,AppleUpdateRegistryProblemFlagsCallback)
#endif

// string to hex
USAGE_AND_PAGE AppleKeyboardString2Hex(__in PWCH pString,__in ULONG ulLength)
{
	ULONG ulRet = 0;

	for(ULONG i = 0; pString[i] && i < ulLength; i ++)
	{
		WCHAR ch = pString[i];

		if(ch >= L'0' && ch <= L'9')
			ulRet = ulRet * 16 + ch - L'0';
		else if(ch >= L'a' && ch <= L'f')
			ulRet = ulRet * 16 + ch - L'a' + 10;
		else if(ch >= L'A' && ch <= L'F')
			ulRet = ulRet * 16 + ch - L'A' + 10;
	}

	return *static_cast<PUSAGE_AND_PAGE>(static_cast<PVOID>(&ulRet));
}

// load usage list
VOID AppleKeyboardLoadUsageMappingList(__in CDeviceExtension* pDevExt)
{
	HANDLE hDriverKey = 0;
	HANDLE hMappingKey = 0;
	CUsageMapping* pLastMapping = 0;

	__try
	{
		if(!NT_SUCCESS(IoOpenDeviceRegistryKey(pDevExt->m_pPhysicalDeviceObject,PLUGPLAY_REGKEY_DRIVER,KEY_ALL_ACCESS,&hDriverKey)))
			__leave;

		UNICODE_STRING name;
		RtlInitUnicodeString(&name,L"UsageAndPageMappings");

		OBJECT_ATTRIBUTES oa;
		RtlZeroMemory(&oa,sizeof(oa));
		oa.RootDirectory = hDriverKey;
		oa.Length = sizeof(oa);
		oa.Attributes = OBJ_CASE_INSENSITIVE;
		oa.ObjectName = &name;

		if(!NT_SUCCESS(ZwOpenKey(&hMappingKey,KEY_ALL_ACCESS,&oa)))
			__leave;

		for(ULONG i = 0; ; i ++)
		{
			ULONG length;
			UCHAR buffer[sizeof(KEY_VALUE_FULL_INFORMATION) + sizeof(WCHAR)*(sizeof(USAGE_AND_PAGE) * 4 + 2)] = {0};
			if(!NT_SUCCESS(ZwEnumerateValueKey(hMappingKey,i,KeyValueFullInformation,buffer,sizeof(buffer),&length)))
				break;

			PKEY_VALUE_FULL_INFORMATION pInfo = static_cast<PKEY_VALUE_FULL_INFORMATION>(static_cast<PVOID>(buffer));
			ASSERT(pInfo->DataLength == (sizeof(USAGE_AND_PAGE) * 2 + 1) * sizeof(WCHAR));
			ASSERT(pInfo->Type == REG_SZ);
			ASSERT(pInfo->NameLength <= (sizeof(USAGE_AND_PAGE) * 2 + 1) * sizeof(WCHAR));

			USAGE_AND_PAGE ulFrom = AppleKeyboardString2Hex(pInfo->Name,pInfo->NameLength / sizeof(WCHAR));
			USAGE_AND_PAGE ulTo = AppleKeyboardString2Hex(static_cast<PWCH>(static_cast<PVOID>(buffer + pInfo->DataOffset)),pInfo->DataLength / sizeof(WCHAR));

			CUsageMapping* pCurMapping = static_cast<CUsageMapping*>(ExAllocatePoolWithTag(NonPagedPool,sizeof(CUsageMapping),'Appl'));
			if(!pCurMapping)
				__leave;

			pCurMapping->m_ulFrom.Usage = ulFrom.Usage;
			pCurMapping->m_ulFrom.UsagePage = ulFrom.UsagePage;
			pCurMapping->m_ulTo.Usage = ulTo.Usage;
			pCurMapping->m_ulTo.UsagePage = ulTo.UsagePage;
			pCurMapping->m_pNextMapping = pLastMapping;
			pLastMapping = pCurMapping;
		}
	}
	__finally
	{
		if(hMappingKey)
			ZwClose(hMappingKey);

		if(hDriverKey)
			ZwClose(hDriverKey);
	}

	KIRQL saveIrql;
	KeAcquireSpinLock(&pDevExt->m_spinLockUsageMappings,&saveIrql);
	pDevExt->m_pUsageIdMappings = pLastMapping;
	KeReleaseSpinLock(&pDevExt->m_spinLockUsageMappings,saveIrql);
}

// map usage
VOID AppleKeyboardMapUsage(__in CDeviceExtension* pDevExt,__in PUSAGE_AND_PAGE pUsageAndPage)
{
	KIRQL saveIrql;
	KeAcquireSpinLock(&pDevExt->m_spinLockUsageMappings,&saveIrql);

	CUsageMapping* pMap = pDevExt->m_pUsageIdMappings;
	while(pMap)
	{
		if(pMap->m_ulFrom.UsagePage == pUsageAndPage->UsagePage && pMap->m_ulFrom.Usage == pUsageAndPage->Usage)
		{
			pUsageAndPage->Usage = pMap->m_ulTo.Usage;
			pUsageAndPage->UsagePage = pMap->m_ulTo.UsagePage;
			break;
		}

		pMap = pMap->m_pNextMapping;
	}

	KeReleaseSpinLock(&pDevExt->m_spinLockUsageMappings,saveIrql);
}

// free usageid mappings
VOID AppleKeyboardFreeUsageMappingList(__in CDeviceExtension* pDevExt)
{
	KIRQL saveIrql = 0;
	KeAcquireSpinLock(&pDevExt->m_spinLockUsageMappings,&saveIrql);
	PVOID pBuffer = pDevExt->m_pUsageIdMappings;
	pDevExt->m_pUsageIdMappings = 0;
	KeReleaseSpinLock(&pDevExt->m_spinLockUsageMappings,saveIrql);

	if(pBuffer)
		ExFreePool(pBuffer);
}

// call hid class
NTSTATUS AppleKeyboardCallHidClass(__in CDeviceExtension* pDevExt,__in ULONG ulIoControlCode,__in PVOID pInputBuffer,__in ULONG ulInputLength,
								   __out PVOID pOutBuffer,__in ULONG ulOutputLength)
{
	PAGED_CODE();

	KEVENT eventComplete;
	KeInitializeEvent(&eventComplete,NotificationEvent,FALSE);

	IO_STATUS_BLOCK statusBlock;
	PIRP pIrp = IoBuildDeviceIoControlRequest(ulIoControlCode,pDevExt->m_pLowerDevice,pInputBuffer,ulInputLength,pOutBuffer,ulOutputLength,
											  FALSE,&eventComplete,&statusBlock);

	if(!pIrp)
		return STATUS_NO_MEMORY;

	NTSTATUS status = IoCallDriver(pDevExt->m_pLowerDevice,pIrp);
	if(status == STATUS_PENDING)
		status = KeWaitForSingleObject(&eventComplete,Executive,KernelMode,FALSE,0);

	if(NT_SUCCESS(status))
		status = statusBlock.Status;

	return status;
}

// query reg key
NTSTATUS AppleKeyboardQueryDeviceKey(__in HANDLE hKey,__in PWCH pValueName,__in __out PVOID pDataBuffer,__in ULONG ulDataLength)
{
	NTSTATUS statusRet = STATUS_SUCCESS;
	PVOID pBuffer = 0;

	__try
	{
		UNICODE_STRING valueName;
		RtlInitUnicodeString(&valueName,pValueName);
		
		pBuffer = ExAllocatePoolWithTag(PagedPool,ulDataLength + sizeof(KEY_VALUE_FULL_INFORMATION),'Appl');
		if(!pBuffer)
		{
			statusRet = STATUS_NO_MEMORY;
			__leave;
		}
		
		ULONG length = 0;
		statusRet = ZwQueryValueKey(hKey,&valueName,KeyValueFullInformation,pBuffer,ulDataLength + sizeof(KEY_VALUE_FULL_INFORMATION),&length);
		if(!NT_SUCCESS(statusRet))
			__leave;

		PKEY_VALUE_FULL_INFORMATION pInfo = static_cast<PKEY_VALUE_FULL_INFORMATION>(pBuffer);
		ASSERT(ulDataLength == pInfo->DataLength);

		RtlCopyMemory(pDataBuffer,pInfo->DataOffset + static_cast<PUCHAR>(pBuffer),pInfo->DataLength);
	}
	__finally
	{
		if(pBuffer)
			ExFreePool(pBuffer);
	}

	return statusRet;
}

// update callback
VOID AppleUpdateRegistryProblemFlagsCallback(__in PDEVICE_OBJECT pDeviceObject,__in_opt PVOID pContext)
{
	PAGED_CODE();

	CDeviceExtension* pDevExt = static_cast<CDeviceExtension*>(pDeviceObject->DeviceExtension);

	HANDLE hKey = 0;
	if(NT_SUCCESS(IoOpenDeviceRegistryKey(pDevExt->m_pPhysicalDeviceObject,PLUGPLAY_REGKEY_DEVICE,KEY_ALL_ACCESS,&hKey)))
	{
		UNICODE_STRING valueName;
		RtlInitUnicodeString(&valueName,L"ProblemFlags");
		ZwSetValueKey(hKey,&valueName,0,REG_DWORD,&pDevExt->m_ulProblemFlags,sizeof(pDevExt->m_ulProblemFlags));
		ZwClose(hKey);
	}

	IoReleaseRemoveLock(&pDevExt->m_removeLock,&AppleUpdateRegistryProblemFlags);
	IoFreeWorkItem(static_cast<PIO_WORKITEM>(pContext));
}

// set problem flags
VOID AppleUpdateRegistryProblemFlags(__in CDeviceExtension* pDevExt,__in PUSAGE_AND_PAGE pInputUsage)
{
	if(!NT_SUCCESS(IoAcquireRemoveLock(&pDevExt->m_removeLock,&AppleUpdateRegistryProblemFlags)))
		return;

	BOOLEAN bReleaseRemoveLock = TRUE;
	__try
	{
		PIO_WORKITEM pIoWorkItem = IoAllocateWorkItem(pDevExt->m_pFunctionDeviceObject);
		if(!pIoWorkItem)
			__leave;
		
		bReleaseRemoveLock = FALSE;
		if(KeGetCurrentIrql() == PASSIVE_LEVEL)
			AppleUpdateRegistryProblemFlagsCallback(pDevExt->m_pFunctionDeviceObject,pIoWorkItem);
		else
			IoQueueWorkItem(pIoWorkItem,&AppleUpdateRegistryProblemFlagsCallback,DelayedWorkQueue,pIoWorkItem);
	}
	__finally
	{
		if(bReleaseRemoveLock)
			IoReleaseRemoveLock(&pDevExt->m_removeLock,&AppleUpdateRegistryProblemFlags);
	}
}